import json
import random

from tqdm import tqdm

import config
from lib import cut

flags = [0, 0, 0, 1]  # 1/5的位置做测试数据集

# 闲聊预料
xiaohuangji_path = r"E:\shenhaitaoPyCode\chat_service\corpus\classify\orgin_corpus\小黄鸡未分词.conv"
# 问答语料
byhand_path = r"E:\shenhaitaoPyCode\chat_service\corpus\classify\orgin_corpus\手动构造的问题.json"
crawled_path = r"E:\shenhaitaoPyCode\chat_service\corpus\classify\orgin_corpus\爬虫抓取的问题.csv"


def keywords_in_line(line):
    keywords_list = ["传智播客", "传智", "黑马程序员", "黑马", "python", "人工智能", "c语言", "c++", "java", "javaee", "前端", "移动开发", "ui",
                     "ue", "大数据", "软件测试", "php", "h5", "产品经理", "linux", "运维", "go语言", "区块链", "影视制作", "pmp", "项目管理",
                     "新媒体", "小程序", "前端"]
    for word in line:
        if word in keywords_list:
            return True
        else:
            return False


def process_xiaohuangji(f_train, f_test):
    """处理小黄鸡语料"""
    train_num = 0
    test_num = 0
    for line in tqdm(open(xiaohuangji_path, encoding="utf-8").readlines(), desc="小黄鸡"):
        # TODO 句子长度为1，考虑删除
        if line.startswith("E"):
            flag = 0
            continue
        elif line.startswith("M"):
            if flag == 0:
                line = line[1:].strip()
                flag = 1
            else:
                continue  # 不需要第二个出现M开头的句子

        line_cuted = cut(line)
        if not keywords_in_line(line_cuted):
            line_cuted = " ".join(line_cuted) + "\t" + "__lable__chat"
            if random.choice(flags) == 0:
                train_num += 1
                f_train.write(line_cuted + "\n")
            else:
                test_num += 1
                f_test.write(line_cuted + "\n")
    return train_num, test_num


def process_byhand_data(f_train, f_test):
    train_num = 0
    test_num = 0
    total_lines = json.loads(open(byhand_path, encoding="utf-8").read())
    for key in total_lines:
        for lines in tqdm(total_lines[key], desc="手动构造问题"):
            for line in lines:
                if "校区" in line:
                    continue
                line_cuted = cut(line)

                line_cuted = " ".join(line_cuted) + "\t" + "__lable__QA"
                if random.choice(flags) == 0:
                    train_num += 1
                    f_train.write(line_cuted + "\n")
                else:
                    test_num += 1
                    f_test.write(line_cuted + "\n")
    return train_num, test_num


def prcess_crawled_data(f_train, f_test):
    train_num = 0
    test_num = 0
    """处理抓取的问题数据"""
    for line in tqdm(open(crawled_path, encoding="utf-8").readlines(), desc="爬虫抓取问题"):
        line = line[:len(line) - 1]
        line_cuted = cut(line)
        line_cuted = " ".join(line_cuted) + "\t" + "__lable__QA"
        if random.choice(flags) == 0:
            train_num += 1
            f_train.write(line_cuted + "\n")
        else:
            test_num += 1
            f_test.write(line_cuted + "\n")
    return train_num, test_num


def process():
    f_train = open(config.classify_corpus_train_path, "a", encoding="utf-8")
    f_test = open(config.classify_corpus_test_path, "a", encoding="utf-8")
    # 1.处理小黄鸡
    num_chat_train, num_chat_test = process_xiaohuangji(f_train, f_test)
    # 2.处理手动构造的问题
    num_qa_train, num_qa_test = process_byhand_data(f_train, f_test)
    # 3.处理爬虫抓取的问题
    _a, _b = prcess_crawled_data(f_train, f_test)
    num_qa_train += _a
    num_qa_test += _b
    f_train.close()
    f_test.close()
    print("训练集和测试集：", num_chat_train + num_qa_train, num_chat_test + num_qa_test)
    print("QA语料和chat语料：", num_qa_train + num_qa_test, num_chat_train + num_chat_test)
